// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/substrait/proto/substrait/algebra.pb.h>
#include <pollux/substrait/proto/substrait/capabilities.pb.h>
#include <pollux/substrait/proto/substrait/extensions/extensions.pb.h>
#include <pollux/substrait/proto/substrait/function.pb.h>
#include <pollux/substrait/proto/substrait/parameterized_types.pb.h>
#include <pollux/substrait/proto/substrait/plan.pb.h>
#include <pollux/substrait/proto/substrait/type.pb.h>
#include <pollux/substrait/proto/substrait/type_expressions.pb.h>

#include <pollux/type/type.h>

namespace kumo::pollux::substrait {
    /// This class contains some common functions used to parse Substrait
    /// components, and convert them into recognizable representations.
    class SubstraitParser {
    public:
        /// Parse Substrait NamedStruct.
        std::vector<TypePtr> parseNamedStruct(
            const ::substrait::NamedStruct &namedStruct);

        /// Parse Substrait Type.
        TypePtr parseType(const ::substrait::Type &substraitType);

        /// Parse Substrait ReferenceSegment.
        int32_t parseReferenceSegment(
            const ::substrait::Expression::ReferenceSegment &refSegment);

        /// Make names in the format of {prefix}_{index}.
        std::vector<std::string> makeNames(const std::string &prefix, int size);

        /// Make node name in the format of n{nodeId}_{colIdx}.
        std::string makeNodeName(int nodeId, int colIdx);

        /// Get the column index from a node name in the format of
        /// n{nodeId}_{colIdx}.
        int getIdxFromNodeName(const std::string &nodeName);

        /// Find the Substrait function name according to the function id
        /// from a pre-constructed function map. The function specification can be
        /// a simple name or a compound name. The compound name format is:
        /// <function name>:<short_arg_type0>_<short_arg_type1>_..._<short_arg_typeN>.
        /// Currently, the input types in the function specification are not used. But
        /// in the future, they should be used for the validation according the
        /// specifications in Substrait yaml files.
        const std::string &findFunctionSpec(
            const melon::F14FastMap<uint64_t, std::string> &functionMap,
            uint64_t id) const;

        /// This function is used get the types from the compound name.
        static std::vector<std::string> getSubFunctionTypes(
            const std::string &substraitFunction);

        /// Find the Pollux function name according to the function id
        /// from a pre-constructed function map.
        std::string findPolluxFunction(
            const melon::F14FastMap<uint64_t, std::string> &functionMap,
            uint64_t id) const;

        /// Map the Substrait function keyword into Pollux function keyword.
        std::string mapToPolluxFunction(const std::string &substraitFunction) const;

        /// Get input types from Substrait function signature.
        static std::vector<TypePtr> getInputTypes(const std::string &signature);

    private:
        /// A map used for mapping Substrait function keywords into Pollux functions'
        /// keywords. Key: the Substrait function keyword, Value: the Pollux function
        /// keyword. For those functions with different names in Substrait and Pollux,
        /// a mapping relation should be added here.
        melon::F14FastMap<std::string, std::string> substraitPolluxFunctionMap_ = {
            {"add", "plus"},
            {"subtract", "minus"},
            {"modulus", "mod"},
            {"not_equal", "neq"},
            {"equal", "eq"}
        };
    };
} // namespace kumo::pollux::substrait
